import requests
from bs4 import BeautifulSoup
import re
import json

headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36 Edg/120.0.0.0",
    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
    "Accept-Encoding": "gzip, deflate, br"
}
base_url = "https://www.stats.gov.cn/sj/tjbz/tjyqhdmhcxhfdm/2023/"


# 返回下一级URL的基地址
def getNextLevelUrlPrefix(currentUrl):
    return re.sub(r"\d+.html$", "", currentUrl)


# 请求国家统计局网站获取HTML代码
def getHtmlRegionTr(url, class_name):
    response = requests.get(url, headers=headers)
    response.encoding = response.status_code

    html = response.text
    html_soup = BeautifulSoup(html, "html.parser")

    # 找到所有的行政区tr
    all_tr = html_soup.findAll("tr", attrs={"class": class_name})
    return all_tr


# 根据街道url获取其下所有的社区级行政区划（五级）
def getVillages(url):
    village_list = []
    all_village_tr = getHtmlRegionTr(url, "villagetr")

    for tr in all_village_tr:
        village_td_list = tr.findAll("td")
        print(f"\t\t\t\t{village_td_list[0].string}{village_td_list[2].string}")
        village_list.append({
            "code": village_td_list[0].string,
            "name": village_td_list[2].string,
            "typeCode": village_td_list[1].string
        })

    return village_list


# 根据区县url获取其下所有的街道级行政区划（四级）
def getTowns(url):
    town_list = []
    all_town_tr = getHtmlRegionTr(url, "towntr")

    for tr in all_town_tr:
        town_td_list = tr.findAll("td")
        print(f"\t\t\t{town_td_list[0].string}{town_td_list[1].string}")
        town_dict = {
            "code": town_td_list[0].string,
            "name": town_td_list[1].string
        }
        # 找td下面有没有a标签，如果有a标签，则继续请求下级行政区划
        code_td_a = town_td_list[0].find("a")
        if (code_td_a):
            try:
                villages = getVillages(f"{getNextLevelUrlPrefix(url)}{code_td_a['href']}")
                town_dict["children"] = villages
            except Exception as e:
                print(f"\t\t\t{town_td_list[1].string}没有下一级行政区")
                # print(f"\t\t\t\t出现异常了:{e}，异常类型：{type(e)}")

        town_list.append(town_dict)

    return town_list


# 根据市url获取其下所有的县级行政区（三级）
def getCounties(url):
    county_list = []
    all_county_tr = getHtmlRegionTr(url, "countytr")

    for tr in all_county_tr:
        county_td_list = tr.findAll("td")
        print("\t\t", county_td_list[0].string, county_td_list[1].string)
        county_dict = {
            "code": county_td_list[0].string,
            "name": county_td_list[1].string
        }
        # 找td下面有没有a标签，如果有a标签，则继续请求下级行政区划
        code_td_a = county_td_list[0].find("a")
        if (code_td_a):
            try:
                towns = getTowns(f"{getNextLevelUrlPrefix(url)}{code_td_a['href']}")
                county_dict["children"] = towns
            except Exception as e:
                print(f"\t\t{county_td_list[1].string}没有下一级行政区")
                # print(f"\t\t\t出现异常了:{e}，异常类型：{type(e)}")

        county_list.append(county_dict)

    return county_list


# 根据省份url获取其下所有的市级行政区（二级）
def getCities(url):
    city_list = []
    all_city_tr = getHtmlRegionTr(url, "citytr")

    for tr in all_city_tr:
        # 找到某省所有市级行政区划代码和名称，下标0是代码，1是名称
        city_a_list = tr.findAll("a")
        print("\t", city_a_list[0].string, city_a_list[1].string)
        city_dict = {
            "code": city_a_list[0].string,
            "name": city_a_list[1].string
        }
        try:
            # 获取该市下所有的县级行政区
            counties = getCounties(f"{getNextLevelUrlPrefix(url)}{city_a_list[0]['href']}")
            city_dict["children"] = counties
        except Exception as e:
            print(f"\t\t{city_a_list[1].string}没有下一级行政区")
            # print(f"\t\t出现异常了:{e}，异常类型：{type(e)}")

        city_list.append(city_dict)

    return city_list


# 获取所有的省份（一级）
def getAllProvince():
    provice_list = []
    all_provinces_tr = getHtmlRegionTr(base_url + "index.html", "provincetr")

    for tr in all_provinces_tr:
        # 找到所有的省份名称及其对应的链接
        tr_a_list = tr.findAll("a")
        for a in tr_a_list:
            print(a.text + ":")
            # 获取当前省下面的所有市级行政区
            cities = getCities(base_url + a['href'])
            provice_list.append({"name": a.text, "children": cities})

    return provice_list


provice_list = getAllProvince()
json_str = json.dumps(provice_list, ensure_ascii=False)
print(json_str)
